//===--- ARCBBState.h -------------------------------------------*- C++ -*-===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#ifndef POLARPHP_PIL_OPTIMIZER_INTERNAL_ARC_ARCBBSTATE_H
#define POLARPHP_PIL_OPTIMIZER_INTERNAL_ARC_ARCBBSTATE_H

#include "polarphp/pil/optimizer/internal/arc/GlobalARCSequenceDataflow.h"

namespace polar {

/// Per-BasicBlock state.
class ARCSequenceDataflowEvaluator::ARCBBState {
public:
   using TopDownMapTy = SmallBlotMapVector<PILValue, TopDownRefCountState, 4>;
   using BottomUpMapTy = SmallBlotMapVector<PILValue, BottomUpRefCountState, 4>;

private:
   /// The basic block that this bbstate corresponds to.
   PILBasicBlock *BB;

   /// The top-down traversal uses this to record information known about a
   /// pointer at the bottom of each block.
   TopDownMapTy PtrToTopDownState;

   /// The bottom-up traversal uses this to record information known about a
   /// pointer at the top of each block.
   BottomUpMapTy PtrToBottomUpState;

   /// Is this a BB that is a trap?
   bool IsTrapBB;

public:
   ARCBBState() : BB() {}
   ARCBBState(PILBasicBlock *BB) : BB(BB) {}

   void init(PILBasicBlock *NewBB, bool NewIsTrapBB) {
      assert(NewBB && "Cannot set NewBB to a nullptr.");
      BB = NewBB;
      IsTrapBB = NewIsTrapBB;
   }

   /// Is this BB a BB that fits the canonical form of a trap?
   ///
   /// The canonical form of a trap is:
   ///
   ///   builtin "int_trap"() : $()
   ///   unreachable
   ///
   /// This cannot have any uses of reference counted values since the frontend
   /// just leaks at that point.
   bool isTrapBB() const { return IsTrapBB; }

   /// Top Down Iterators
   using topdown_iterator = TopDownMapTy::iterator;
   using topdown_const_iterator = TopDownMapTy::const_iterator;
   topdown_iterator topdown_begin() { return PtrToTopDownState.begin(); }
   topdown_iterator topdown_end() { return PtrToTopDownState.end(); }
   topdown_const_iterator topdown_begin() const {
      return PtrToTopDownState.begin();
   }
   topdown_const_iterator topdown_end() const { return PtrToTopDownState.end(); }
   iterator_range<topdown_iterator> getTopDownStates() {
      return make_range(topdown_begin(), topdown_end());
   }

   /// Bottom up iteration.
   using bottomup_iterator = BottomUpMapTy::iterator;
   using bottomup_const_iterator = BottomUpMapTy::const_iterator;
   bottomup_iterator bottomup_begin() { return PtrToBottomUpState.begin(); }
   bottomup_iterator bottomup_end() { return PtrToBottomUpState.end(); }
   bottomup_const_iterator bottomup_begin() const {
      return PtrToBottomUpState.begin();
   }
   bottomup_const_iterator bottomup_end() const {
      return PtrToBottomUpState.end();
   }
   iterator_range<bottomup_iterator> getBottomupStates() {
      return make_range(bottomup_begin(), bottomup_end());
   }

   /// Attempt to find the PtrState object describing the top down state for
   /// pointer Arg. Return a new initialized PtrState describing the top down
   /// state for Arg if we do not find one.
   TopDownRefCountState &getTopDownRefCountState(PILValue Ptr) {
      return PtrToTopDownState[Ptr];
   }

   /// Attempt to find the PtrState object describing the bottom up state for
   /// pointer Arg. Return a new initialized PtrState describing the bottom up
   /// state for Arg if we do not find one.
   BottomUpRefCountState &getBottomUpRefCountState(PILValue Ptr) {
      return PtrToBottomUpState[Ptr];
   }

   /// Blot \p Ptr.
   void clearBottomUpRefCountState(PILValue Ptr) {
      PtrToBottomUpState.erase(Ptr);
   }

   /// Blot \p Ptr.
   void clearTopDownRefCountState(PILValue Ptr) { PtrToTopDownState.erase(Ptr); }

   void clearTopDownState() { PtrToTopDownState.clear(); }
   void clearBottomUpState() { PtrToBottomUpState.clear(); }

   /// Clear both the bottom up *AND* top down state.
   void clear() {
      clearTopDownState();
      clearBottomUpState();
   }

   /// Returns a reference to the basic block that we are tracking.
   PILBasicBlock &getBB() const { return *BB; }

   /// Merge in the state of the successor basic block. This is currently a stub.
   void mergeSuccBottomUp(ARCBBState &SuccBB);

   /// Initialize this BB with the state of the successor basic block. This is
   /// called on a basic block's state and then any other successors states are
   /// merged in. This is currently a stub.
   void initSuccBottomUp(ARCBBState &SuccBB);

   /// Merge in the state of the predecessor basic block. This is currently a
   /// stub.
   void mergePredTopDown(ARCBBState &PredBB);

   /// Initialize the state for this BB with the state of its predecessor
   /// BB. Used to create an initial state before we merge in other
   /// predecessors. This is currently a stub.
   void initPredTopDown(ARCBBState &PredBB);
};

class ARCSequenceDataflowEvaluator::ARCBBStateInfoHandle {
   friend ARCSequenceDataflowEvaluator::ARCBBStateInfo;

   PILBasicBlock *BB;
   ARCBBState &BBState;
   NullablePtr<llvm::SmallPtrSet<PILBasicBlock *, 4>> BackedgeMap;
   unsigned ID;

   ARCBBStateInfoHandle(PILBasicBlock *BB, unsigned ID, ARCBBState &BBState)
      : BB(BB), BBState(BBState), BackedgeMap(), ID(ID) {}
   ARCBBStateInfoHandle(PILBasicBlock *BB, unsigned ID, ARCBBState &BBState,
                        llvm::SmallPtrSet<PILBasicBlock *, 4> &BackedgeMap)
      : BB(BB), BBState(BBState), BackedgeMap(&BackedgeMap), ID(ID) {}

public:
   ARCBBStateInfoHandle(const ARCBBStateInfoHandle &) = default;
   ARCBBStateInfoHandle(ARCBBStateInfoHandle &&) = default;
   ~ARCBBStateInfoHandle() = default;

   PILBasicBlock *getBB() const { return BB; }
   unsigned getID() const { return ID; }
   ARCBBState &getState() { return BBState; }
   bool isBackedge(PILBasicBlock *BB) const {
      if (BackedgeMap.isNull())
         return false;
      return BackedgeMap.get()->count(BB);
   }
};

class ARCSequenceDataflowEvaluator::ARCBBStateInfo {
   /// A map from BB -> BBID. A BB's BBID is its RPOT number.
   llvm::DenseMap<PILBasicBlock *, unsigned> BBToBBIDMap;

   /// Map from a BBID to BB's bottom up dataflow state. Meant to be used in
   /// conjunction with BBToBBIDMap.
   std::vector<ARCBBState> BBIDToBottomUpBBStateMap;

   /// Map from a BBID to BB's top down dataflow state. Meant to be used in
   /// conjunction with BBToBBIDMap.
   std::vector<ARCBBState> BBIDToTopDownBBStateMap;

   /// A map mapping the head to a tail of a backedge. We only compute this once
   /// in the lifetime of this class.
   llvm::DenseMap<PILBasicBlock *, llvm::SmallPtrSet<PILBasicBlock *, 4>>
   BackedgeMap;

public:
   ARCBBStateInfo(PILFunction *F, PostOrderAnalysis *POTA,
                  ProgramTerminationFunctionInfo *PTFI);

   llvm::Optional<ARCBBStateInfoHandle> getBottomUpBBHandle(PILBasicBlock *BB);
   llvm::Optional<ARCBBStateInfoHandle> getTopDownBBHandle(PILBasicBlock *BB);

   void clear();

private:
   llvm::Optional<unsigned> getBBID(PILBasicBlock *BB) const;
};

} // end swift namespace

#endif // POLARPHP_PIL_OPTIMIZER_INTERNAL_ARC_ARCBBSTATE_H
